`timescale 1ns / 1ps
//////////////////////////////////////////////////////////////////////////////////
// Company: CBICR, Tsinghua Univ.
// Engineer: Hongyi Li & Kimi & Deepseek
// 
// Create Date: 2025/03/16 20:05
// Design Name: 
// Module Name: Currying ALU in CompAir
// Project Name: 
// Target Devices: 
// Tool Versions: 
// Description: 
// 
// Dependencies: 
// 
// Revision:
// Revision 0.01 - File Created
// Additional Comments:
// 
//////////////////////////////////////////////////////////////////////////////////

module CurryALU
(
    input                     clk, rst_n,
    input  [15:0]             i_data,
    input  [ 1:0]             i_op,
    input                     i_iter_tag,
    input                     i_wr_force,
    input                     i_wr_result,
    input                     i_rd_acc,
    output [15:0]             o_data
);

reg         last_iter_reg;
reg  [15:0] acc_data_reg;
reg  [15:0] acc_iter_reg;
reg  [ 1:0] acc_op_reg;


wire [15:0]  data_a;
wire [15:0]  data_b;

wire sign_b_inpt;
wire sign_b_iter;

assign data_a = acc_data_reg;
assign sign_b_inpt = (i_op == 2'b01) ? (~i_data[15]) : i_data[15];
assign sign_b_iter = (i_op == 2'b01) ? (~acc_iter_reg[15]) : acc_iter_reg[15];

assign data_b = last_iter_reg ? {sign_b_iter , acc_iter_reg[14:0]} : {sign_b_inpt , i_data[14:0]};

wire [15:0] add_res;
wire [15:0] mul_res;
wire [15:0] div_res;
wire [15:0] alu_res;

// Add & Sub
bf16_adder add (
    .a(data_a),
    .b(data_b),
    .sum(add_res)
);

// Mul
bf16_multiplier mul (
    .a(data_a),
    .b(data_b),
    .product(mul_res)
);

// Div
bf16_divider div (
    .a(data_a),
    .b(data_b),
    .quotient(div_res)
);

assign alu_res = (i_op == 2'b00) ? add_res : 
                 (i_op == 2'b01) ? add_res : 
                 (i_op == 2'b10) ? mul_res :
                 (i_op == 2'b11) ? div_res :
                 16'b0;

// Input-Related Reg Update
always @(posedge clk) begin

    // last_iter_tag
    if (!rst_n) 
        last_iter_reg <= 0;
    else
        last_iter_reg <= i_iter_tag;

    // acc_data_reg
    if (!rst_n) 
        acc_data_reg <= 0;
    else if (i_wr_force)
        acc_data_reg <= i_data;
    else if (last_iter_reg || i_wr_result)
        acc_data_reg <= alu_res;
    else
        acc_data_reg <= acc_data_reg;
    
    // acc_iter_reg & acc_op_reg
    if (!rst_n) begin
        acc_iter_reg <= 0;
        acc_op_reg <= 0;
    end else if (i_wr_force && i_iter_tag) begin
        acc_iter_reg <= i_data;
        acc_op_reg <= i_op;
    end else begin
        acc_iter_reg <= acc_iter_reg;
        acc_op_reg <= acc_op_reg;
    end
end

assign o_data = acc_data_reg;

endmodule